package io.github.melin.sqlflow.metadata;

import io.github.melin.sqlflow.DefaultTraversalVisitor;
import io.github.melin.sqlflow.tree.SelectItem;
import io.github.melin.sqlflow.tree.SingleColumn;
import io.github.melin.sqlflow.tree.With;
import io.github.melin.sqlflow.tree.WithQuery;
import io.github.melin.sqlflow.tree.expression.*;
import io.github.melin.sqlflow.tree.join.Join;
import io.github.melin.sqlflow.tree.join.JoinCriteria;
import io.github.melin.sqlflow.tree.join.JoinOn;
import io.github.melin.sqlflow.tree.join.JoinUsing;
import io.github.melin.sqlflow.tree.relation.*;
import io.github.melin.sqlflow.tree.statement.CreateTableAsSelect;
import io.github.melin.sqlflow.tree.statement.Insert;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.stream.Collectors;

/**
 * SQL 元数据提取器，用于分析 SQL 语句中的表和字段信息
 */
public class SqlMetadataExtractor extends DefaultTraversalVisitor<Void> {
    // 使用ThreadLocal存储线程私有的状态
    private static final ThreadLocal<SqlMetadataContext> context = ThreadLocal.withInitial(SqlMetadataContext::new);

    // 内部类用于存储线程私有的上下文
    private static class SqlMetadataContext {
        // 记录当前正在处理的表名（包括表别名）
        private String currentTable = null;
        private final String defaultSchema="default";

        // 存储源表信息（SELECT FROM 中的表）- 使用ConcurrentHashMap和CopyOnWriteArraySet保持线程安全
        private final Map<String, Set<String>> sourceTablesWithColumns = new ConcurrentHashMap<>();

        // 存储目标表信息（INSERT INTO 的表）
        private SchemaTable targetTable = null;

        // 存储表和列的映射关系 - 使用ConcurrentHashMap和CopyOnWriteArraySet保持线程安全
        private final Map<String, Set<String>> tableColumns = new ConcurrentHashMap<>();

        // 存储表别名到真实表名的映射
        private final Map<String, String> tableAliases = new ConcurrentHashMap<>();

        // 存储最后一个 SELECT 语句的输出列
        private final List<String> lastSelectColumns = new CopyOnWriteArrayList<>();

        // 存储子查询中的表和列信息 - 使用ConcurrentHashMap和CopyOnWriteArraySet保持线程安全
        private final Map<String, Set<String>> subqueryColumns = new ConcurrentHashMap<>();

        // 记录当前处理的子查询别名
        private String currentSubqueryAlias = null;

        // 记录FROM子句中表的顺序，用于确定无表前缀列的归属
        private final List<String> tableOrder = new CopyOnWriteArrayList<>();

        // 获取当前上下文的方法
        public void clear() {
            currentTable = null;
            currentSubqueryAlias = null;
            targetTable = null;
            sourceTablesWithColumns.clear();
            tableColumns.clear();
            tableAliases.clear();
            lastSelectColumns.clear();
            subqueryColumns.clear();
            tableOrder.clear();
        }
    }

    // 在每次使用前清理上下文
    public void clearContext() {
        context.get().clear();
    }

    // 在使用完后移除上下文
    public void removeContext() {
        context.remove();
    }
    @Override
    public Void visitCreateTableAsSelect(CreateTableAsSelect node, Void unused) {
        SqlMetadataContext ctx = context.get();

        // 清空上一次的输出列
        ctx.lastSelectColumns.clear();

        // 提取目标表信息
        String targetTableName = node.getName().toString();

        // 初始化目标表的列集合
        ctx.tableColumns.putIfAbsent(targetTableName, new CopyOnWriteArraySet<>());

        // 处理查询部分
        process(node.getQuery(), unused);

        // 使用查询的输出列作为目标表的列
        Set<String> columnSet = new HashSet<>(ctx.lastSelectColumns);

        // 设置目标表
        ctx.targetTable = getSchemaTable(targetTableName, new ArrayList<>(columnSet),null);
        ctx.tableColumns.get(targetTableName).addAll(columnSet);

        // 将目标表添加到源表集合中，确保它被正确识别
        // 这是为了确保在结果中能够看到这个表
        ctx.sourceTablesWithColumns.put(targetTableName, new CopyOnWriteArraySet<>(columnSet));

        return null;
    }

    @Override
    public Void visitIdentifier(Identifier node, Void unused) {
        // 处理标识符（可能是列名）
        addColumnToCurrentContext(node.getValue());
        return null;
    }

    /**
     * 访问引用表达式节点的方法。
     *
     * @param node 引用表达式节点
     * @param unused 未使用的参数
     * @return 返回值为Void类型，表示不返回任何值
     */
    @Override
    public Void visitDereferenceExpression(DereferenceExpression node, Void unused) {
        SqlMetadataContext ctx = context.get();

        // 处理带表别名的列引用
        if (node.getField().isPresent() && node.getField().get() instanceof Identifier) {
            String columnName = ((Identifier) node.getField().get()).getValue();

            // 从node.getBase()获取表别名，确保正确处理Identifier对象
            String tableAlias;
            if (node.getBase() instanceof Identifier) {
                tableAlias = ((Identifier) node.getBase()).getValue();
            } else {
                tableAlias = node.getBase().toString();
            }

            // 获取实际表名
            String realTableName = ctx.tableAliases.getOrDefault(tableAlias, tableAlias);

            // 如果表不存在，先创建它
            ctx.sourceTablesWithColumns.putIfAbsent(realTableName, new CopyOnWriteArraySet<>());

            // 将列添加到实际表中
            ctx.sourceTablesWithColumns.get(realTableName).add(columnName);
            // 如果是表别名，也将列添加到原表中
            if (ctx.tableAliases.containsKey(tableAlias)) {
                String originalTable = ctx.tableAliases.get(tableAlias);
                ctx.sourceTablesWithColumns.putIfAbsent(originalTable, new CopyOnWriteArraySet<>());
                ctx.sourceTablesWithColumns.get(originalTable).add(columnName);
            }

            // 如果这是一个已知的表别名，也将列添加到别名表中
            if (ctx.sourceTablesWithColumns.containsKey(tableAlias)) {
                ctx.sourceTablesWithColumns.get(tableAlias).add(columnName);
            }
        }
        return null;
    }

    @Override
    public Void visitInsert(Insert node, Void unused) {
        SqlMetadataContext ctx = context.get();

        // 清空上一次的输出列
        ctx.lastSelectColumns.clear();

        // 获取目标表信息
        String targetTableName = node.getTarget().toString();

        // 提取 INSERT 语句中的显式指定的列信息
        Set<String> columnSet = new CopyOnWriteArraySet<>();
        if (node.getColumns().isPresent()) {
            node.getColumns().get().forEach(column ->
                    columnSet.add(column.getValue())
            );
        }

        // 提取分区列信息
        List<String> partitionColumns = new ArrayList<>();
        if (node.getPartitionColumns().isPresent()) {
            node.getPartitionColumns().get().forEach((key, value) ->
                    partitionColumns.add(key)
            );
        }

        // 初始化目标表的列集合
        ctx.tableColumns.putIfAbsent(targetTableName, new CopyOnWriteArraySet<>());

        // 处理 WITH 子句
        if (node.getWith().isPresent()) {
            process(node.getWith().get(), unused);
        }

        // 处理查询部分
        process(node.getQuery(), unused);

        // 如果没有显式指定列，使用 SELECT 语句的输出列
        if (columnSet.isEmpty() && !ctx.lastSelectColumns.isEmpty()) {
            columnSet.addAll(ctx.lastSelectColumns);
        }
        ctx.targetTable =getSchemaTable(targetTableName,new ArrayList<>(columnSet),partitionColumns);
        ctx.tableColumns.get(targetTableName).addAll(columnSet);

        return null;
    }

    @Override
    public Void visitQuerySpecification(QuerySpecification node, Void unused) {
        SqlMetadataContext ctx = context.get();

        // 保存当前的表名，以便处理完子查询后恢复
        String savedCurrentTable = ctx.currentTable;
        String savedSubqueryAlias = ctx.currentSubqueryAlias;

        // 清空之前的输出列
        if (ctx.currentSubqueryAlias == null) {
            ctx.lastSelectColumns.clear();
        } else {
            ctx.subqueryColumns.putIfAbsent(ctx.currentSubqueryAlias, new CopyOnWriteArraySet<>());
        }

        // 先处理 FROM 子句
        if (node.getFrom().isPresent()) {
            process(node.getFrom().get(), unused);
        }

        // 处理 WHERE 子句，可能包含对列的引用
        if (node.getWhere().isPresent()) {
            process(node.getWhere().get(), unused);
        }

        // 处理 GROUP BY 子句
        if (node.getGroupBy().isPresent()) {
            process(node.getGroupBy().get(), unused);
        }

        // 处理 HAVING 子句
        if (node.getHaving().isPresent()) {
            process(node.getHaving().get(), unused);
        }

        // 处理 SELECT 子句，收集输出列
        for (SelectItem item : node.getSelect().getSelectItems()) {
            process(item, unused);

            if (item instanceof SingleColumn) {
                SingleColumn singleColumn = (SingleColumn) item;
                String columnName = getColumnName(singleColumn);
                if (columnName != null) {
                    if (ctx.currentSubqueryAlias == null) {
                        ctx.lastSelectColumns.add(columnName);
                    } else {
                        ctx.subqueryColumns.get(ctx.currentSubqueryAlias).add(columnName);
                    }
                }
            }
        }

        // 恢复当前表名
        ctx.currentTable = savedCurrentTable;
        ctx.currentSubqueryAlias = savedSubqueryAlias;

        return null;
    }

    private String getColumnName(SingleColumn node) {
        // 如果有别名，优先使用别名作为输出列名
        if (node.getAlias().isPresent()) {
            return node.getAlias().get().getValue();
        }

        // 没有别名时的处理
        Expression expression = node.getExpression();
        if (expression instanceof DereferenceExpression) {
            // 对于带表别名的列引用，使用列名部分
            DereferenceExpression dref = (DereferenceExpression) expression;
            if (dref.getField().isPresent() && dref.getField().get() instanceof Identifier) {
                return ((Identifier) dref.getField().get()).getValue();
            }
        } else if (expression instanceof Identifier) {
            // 对于简单列引用，使用标识符值
            return ((Identifier) expression).getValue();
        }

        // 对于其他情况（如函数调用），使用表达式的字符串表示
        return expression.toString();
    }

    @Override
    public Void visitTable(Table node, Void unused) {
        SqlMetadataContext ctx = context.get();

        String tableName = node.getName().toString();
        // 清除引号
        tableName = tableName.replaceAll("^\"|\"$", "").replaceAll("^`|`$", "");
        ctx.currentTable = tableName;

        // 初始化表的列集合
        if (!ctx.sourceTablesWithColumns.containsKey(tableName)) {
            ctx.sourceTablesWithColumns.put(tableName, new CopyOnWriteArraySet<>());
        }
        if (!ctx.tableColumns.containsKey(tableName)) {
            ctx.tableColumns.put(tableName, new CopyOnWriteArraySet<>());
        }

        // 将表添加到表顺序列表中，用于确定无表前缀列的归属
        if (!ctx.tableOrder.contains(tableName)) {
            ctx.tableOrder.add(tableName);
        }

        return null;
    }
    @Override
    public Void visitAliasedRelation(AliasedRelation node, Void unused) {
        SqlMetadataContext ctx = context.get();

        // 记录表别名
        String alias = node.getAlias().getValue();
        Relation relation = node.getRelation();

        if (relation instanceof Table) {
            Table table = (Table) relation;
            String tableName = table.getName().toString();
            // 清除引号
            tableName = tableName.replaceAll("^\"|\"$", "").replaceAll("^`|`$", "");
            ctx.tableAliases.put(alias, tableName);

            // 处理表
            process(relation, unused);

            // 将原表的列信息复制到别名表，但不创建别名表的实际记录
            if (ctx.sourceTablesWithColumns.containsKey(tableName)) {
                ctx.sourceTablesWithColumns.putIfAbsent(alias, new CopyOnWriteArraySet<>());
                ctx.sourceTablesWithColumns.get(alias).addAll(ctx.sourceTablesWithColumns.get(tableName));

                // 如果原表在表顺序列表中，也为别名添加相同的顺序位置
                if (ctx.tableOrder.contains(tableName) && !ctx.tableOrder.contains(alias)) {
                    ctx.tableOrder.add(alias);
                }
            }
        }

        // 处理子查询
        super.visitAliasedRelation(node, unused);

        return null;
    }

    @Override
    public Void visitSingleColumn(SingleColumn node, Void unused) {
        SqlMetadataContext ctx = context.get();
        Optional<Identifier> alias = node.getAlias();
        Expression expression = node.getExpression();

        // 获取列名和源表信息
        String sourceColumnName = null;
        String tableAlias = null;
        String realTableName = null;

        // 1. 处理源表列信息
        if (expression instanceof Identifier) {
            // 简单列引用（没有表前缀）
            sourceColumnName = ((Identifier) expression).getValue();

            // 修改后的逻辑：如果没有表前缀，将列添加到FROM子句中第一个表
            if (!ctx.tableOrder.isEmpty()) {
                String firstTable = ctx.tableOrder.get(0);
                if (ctx.sourceTablesWithColumns.containsKey(firstTable)) {
                    ctx.sourceTablesWithColumns.get(firstTable).add(sourceColumnName);
                }
            } else {
                // 如果表顺序列表为空，回退到原有逻辑
                for (Map.Entry<String, Set<String>> entry : ctx.sourceTablesWithColumns.entrySet()) {
                    if (!ctx.subqueryColumns.containsKey(entry.getKey())) {
                        entry.getValue().add(sourceColumnName);
                    }
                }
            }
        } else if (expression instanceof DereferenceExpression) {
            // 带表别名的列引用
            DereferenceExpression dref = (DereferenceExpression) expression;
            if (dref.getField().isPresent() && dref.getField().get() instanceof Identifier) {
                sourceColumnName = ((Identifier) dref.getField().get()).getValue();

                // 改进表别名获取逻辑
                if (dref.getBase() instanceof Identifier) {
                    tableAlias = ((Identifier) dref.getBase()).getValue();
                } else {
                    tableAlias = dref.getBase().toString();
                }

                realTableName = ctx.tableAliases.getOrDefault(tableAlias, tableAlias);
                if (ctx.sourceTablesWithColumns.containsKey(realTableName)) {
                    ctx.sourceTablesWithColumns.get(realTableName).add(sourceColumnName);
                }
            }
        } else if (expression instanceof FunctionCall) {
            // 函数调用
            FunctionCall functionCall = (FunctionCall) expression;
            for (Expression argument : functionCall.getArguments()) {
                if (argument instanceof Identifier) {
                    String argColumnName = ((Identifier) argument).getValue();
                    // 将列添加到所有可能的表中
                    for (Map.Entry<String, Set<String>> entry : ctx.sourceTablesWithColumns.entrySet()) {
                        if (!ctx.subqueryColumns.containsKey(entry.getKey())) {
                            entry.getValue().add(argColumnName);
                        }
                    }
                } else if (argument instanceof DereferenceExpression) {
                    DereferenceExpression dref = (DereferenceExpression) argument;
                    if (dref.getField().isPresent() && dref.getField().get() instanceof Identifier) {
                        String argColumnName = ((Identifier) dref.getField().get()).getValue();
                        String argTableAlias = dref.getBase().toString();
                        String argRealTableName = ctx.tableAliases.getOrDefault(argTableAlias, argTableAlias);
                        if (ctx.sourceTablesWithColumns.containsKey(argRealTableName)) {
                            ctx.sourceTablesWithColumns.get(argRealTableName).add(argColumnName);
                        }
                    }
                }
            }
        }

        // 2. 处理目标表列信息
        String targetColumnName;
        if (alias.isPresent()) {
            // 如果有别名，优先使用别名作为目标列名
            targetColumnName = alias.get().getValue();
        } else if (expression instanceof DereferenceExpression) {
            // 对于带表别名的列引用，使用列名部分
            DereferenceExpression dref = (DereferenceExpression) expression;
            if (dref.getField().isPresent() && dref.getField().get() instanceof Identifier) {
                targetColumnName = ((Identifier) dref.getField().get()).getValue();
            } else {
                targetColumnName = expression.toString();
            }
        } else if (expression instanceof Identifier) {
            // 对于简单列引用，使用标识符值
            targetColumnName = ((Identifier) expression).getValue();
        } else {
            // 对于其他表达式（如函数调用），使用表达式字符串表示
            targetColumnName = expression.toString();
        }

        // 将目标列添加到相应的集合中
        if (isActualColumn(targetColumnName)) {
            if (ctx.currentSubqueryAlias != null) {
                ctx.subqueryColumns.get(ctx.currentSubqueryAlias).add(targetColumnName);
            } else {
                ctx.lastSelectColumns.add(targetColumnName);
            }
        }

        return null;
    }

    // 这个方法已被废弃，使用更专门的方法代替
    private void addColumnToCurrentContext(String columnName) {
        SqlMetadataContext ctx = context.get();
        if (isActualColumn(columnName)) {
            if (ctx.currentSubqueryAlias != null) {
                ctx.subqueryColumns.get(ctx.currentSubqueryAlias).add(columnName);
            } else if (ctx.currentTable != null) {
                ctx.tableColumns.get(ctx.currentTable).add(columnName);
                if (ctx.sourceTablesWithColumns.containsKey(ctx.currentTable)) {
                    ctx.sourceTablesWithColumns.get(ctx.currentTable).add(columnName);
                }
            }
            ctx.lastSelectColumns.add(columnName);
        }
    }

    // 添加源表列 - 添加到源表的列集合中
    private void addSourceColumn(String columnName) {
        SqlMetadataContext ctx = context.get();
        if (isActualColumn(columnName)) {
            // 修改后的逻辑：如果没有表前缀，将列添加到FROM子句中第一个表
            if (!ctx.tableOrder.isEmpty()) {
                String firstTable = ctx.tableOrder.get(0);
                if (ctx.sourceTablesWithColumns.containsKey(firstTable)) {
                    ctx.sourceTablesWithColumns.get(firstTable).add(columnName);
                }
            } else {
                // 如果表顺序列表为空，回退到原有逻辑
                for (Map.Entry<String, Set<String>> entry : ctx.sourceTablesWithColumns.entrySet()) {
                    if (!ctx.subqueryColumns.containsKey(entry.getKey())) {
                        entry.getValue().add(columnName);
                    }
                }
            }
        }
    }

    // 添加目标表列 - 添加到目标表或子查询的列集合中
    private void addTargetColumn(String columnName) {
        SqlMetadataContext ctx = context.get();
        if (isActualColumn(columnName)) {
            if (ctx.currentSubqueryAlias != null) {
                ctx.subqueryColumns.get(ctx.currentSubqueryAlias).add(columnName);
            } else {
                // 对于目标表，总是将列添加到lastSelectColumns中
                ctx.lastSelectColumns.add(columnName);
            }
        }
    }

    // 辅助方法：获取表的真实名称（去除别名）
    private String getRealTableName(String tableNameOrAlias) {
        SqlMetadataContext ctx = context.get();
        return ctx.tableAliases.getOrDefault(tableNameOrAlias, tableNameOrAlias);
    }
    public Void visitWith(With node, Void unused) {
        SqlMetadataContext ctx = context.get();

        // 处理 WITH 子句中的每个查询
        for (WithQuery withQuery : node.getQueries()) {
            String withQueryName = withQuery.getName().getValue();
            ctx.currentSubqueryAlias = withQueryName;
            ctx.subqueryColumns.putIfAbsent(withQueryName, new CopyOnWriteArraySet<>());

            process(withQuery.getQuery(), unused);

            // 将 WITH 查询的列信息添加到源表信息中
            ctx.sourceTablesWithColumns.putIfAbsent(withQueryName, new CopyOnWriteArraySet<>());
            ctx.sourceTablesWithColumns.get(withQueryName).addAll(ctx.subqueryColumns.get(withQueryName));
        }
        ctx.currentSubqueryAlias = null;
        return null;
    }

    public Void visitSubquery(SubqueryExpression node, Void unused) {
        process(node.getQuery(), unused);
        return null;
    }

    @Override
    public Void visitJoin(Join node, Void unused) {
        SqlMetadataContext ctx = context.get();

        // 处理左侧关系
        process(node.getLeft(), unused);

        // 处理右侧关系
        process(node.getRight(), unused);

        // 处理JOIN条件
        if (node.getCriteria().isPresent()) {
            JoinCriteria criteria = node.getCriteria().get();
            if (criteria instanceof JoinOn) {
                process(((JoinOn) criteria).getExpression(), unused);
            } else if (criteria instanceof JoinUsing) {
                // 对于USING子句，我们需要处理列名
                JoinUsing using = (JoinUsing) criteria;
                for (Identifier column : using.getColumns()) {

                    addSourceColumn(column.getValue());
                }
            }
        }

        return null;
    }

    // 移除@Override注解，因为父类DefaultTraversalVisitor中没有此方法
    public Void visitJoinOn(JoinOn node, Void unused) {
        // 处理ON条件中的表达式
        process(node.getExpression(), unused);
        return null;
    }
    @Override
    public Void visitComparisonExpression(ComparisonExpression node, Void unused) {
        // 处理左侧表达式
        if (node.getLeft() instanceof Identifier) {
            // 无表前缀的列引用
            String columnName = ((Identifier) node.getLeft()).getValue();
            addUnprefixedColumn(columnName);
        } else if (node.getLeft() instanceof DereferenceExpression) {
            // 有表前缀的列引用，直接处理，不需要特殊处理
            process(node.getLeft(), unused);
        }

        // 处理右侧表达式
        if (node.getRight() instanceof Identifier) {
            // 无表前缀的列引用
            String columnName = ((Identifier) node.getRight()).getValue();
            addUnprefixedColumn(columnName);
        } else {
            process(node.getRight(), unused);
        }

        return null;
    }

    // 新增辅助方法：处理无表前缀的列
    private void addUnprefixedColumn(String columnName) {
        SqlMetadataContext ctx = context.get();
        if (!ctx.tableOrder.isEmpty()) {
            // 将无表前缀列添加到FROM子句中第一个表
            String firstTable = ctx.tableOrder.get(0);
            if (ctx.sourceTablesWithColumns.containsKey(firstTable)) {
                ctx.sourceTablesWithColumns.get(firstTable).add(columnName);
            }
        }
    }

    /**
     * 获取所有表的信息，包括源表和目标表
     */
    public MetadataService getTableMetadata() {
        SqlMetadataContext ctx = context.get();
        SimpleMetadataService result = new SimpleMetadataService(ctx.defaultSchema);

        // 处理源表信息时，确保只添加实际的表列
        for (Map.Entry<String, Set<String>> entry : ctx.sourceTablesWithColumns.entrySet()) {
            String tableName = entry.getKey();
            // 过滤掉函数调用和复杂表达式，只保留实际的列名
            Set<String> actualColumns = entry.getValue().stream()
                    .filter(this::isActualColumn)
                    .collect(Collectors.toSet());
            result.addTableMetadata(new SchemaTable(tableName, new ArrayList<>(actualColumns)));
        }

        if (ctx.targetTable != null) {
            result.addTableMetadata(ctx.targetTable);
        }
        return result;
    }

    private boolean isActualColumn(String columnName) {
        if (columnName == null || columnName.isEmpty()) {
            return false;
        }

        // 1. 排除包含算术运算符的表达式
        if (columnName.matches(".*[+\\-*/%&|^~<>=!].*")) {
            return false;
        }

        // 2. 如果包含括号，检查是否是完整的函数调用
        if (columnName.contains("(")) {
            // 确保括号是成对的
            int openCount = columnName.length() - columnName.replace("(", "").length();
            int closeCount = columnName.length() - columnName.replace(")", "").length();
            if (openCount != closeCount) {
                return false;
            }

            // 检查是否是有效的函数调用格式
            return columnName.matches("[a-zA-Z_][a-zA-Z0-9_]*\\([^()]*\\)");
        }

        // 3. 基本的列名格式检查 - 允许更宽松的命名规则
        return columnName.matches("[a-zA-Z0-9_]+");
    }

    /**
     * 获取所有表的信息，包括源表和目标表
     */
    public Set<SchemaTable> getTables() {
        SqlMetadataContext ctx = context.get();
        Set<SchemaTable> result = new HashSet<>();
        Map<String, SchemaTable> schemaTableMap = new HashMap<>();

        // 1. 收集表别名到实际表的映射，并合并列信息
        Map<String, Set<String>> realTableColumns = new HashMap<>();
        for (Map.Entry<String, String> aliasEntry : ctx.tableAliases.entrySet()) {
            String alias = aliasEntry.getKey();
            String realTable = aliasEntry.getValue();
            if (ctx.sourceTablesWithColumns.containsKey(alias)) {
                realTableColumns.computeIfAbsent(realTable, k -> new HashSet<>())
                        .addAll(ctx.sourceTablesWithColumns.get(alias));
            }
        }

        // 2. 处理所有实际表（非别名和非子查询）
        for (Map.Entry<String, Set<String>> entry : ctx.sourceTablesWithColumns.entrySet()) {
            String tableName = entry.getKey().trim();

            // 跳过子查询别名和表别名
            if (ctx.subqueryColumns.containsKey(tableName) || ctx.tableAliases.containsKey(tableName)) {
                continue;
            }

            Set<String> allColumns = new HashSet<>(entry.getValue());
            // 合并从别名收集到的列
            if (realTableColumns.containsKey(tableName)) {
                allColumns.addAll(realTableColumns.get(tableName));
            }
            List<String> columns = new ArrayList<>(allColumns);

            // 检查是否已经存在相同的 SchemaTable
            if (schemaTableMap.containsKey(tableName)) {
                // 如果存在，合并列
                SchemaTable existingSchemaTable = schemaTableMap.get(tableName);
                List<String> existingPartitionColumns = existingSchemaTable.getPartitionColumns();
                existingSchemaTable.setColumns(mergeColumns(existingSchemaTable.getColumns(), columns));
                existingSchemaTable.setPartitionColumns(existingPartitionColumns);
            } else {
                // 如果不存在，创建新的 SchemaTable 并添加到集合中
                // 对于源表，我们可能需要从元数据服务获取分区列信息
                // 这里暂时传递 null，因为源表的分区信息通常需要从外部元数据服务获取
                SchemaTable schemaTable = getSchemaTable(tableName, columns, null);
                result.add(schemaTable);
                schemaTableMap.put(tableName, schemaTable);
            }
        }

        // 3. 添加目标表信息
        if (ctx.targetTable != null) {
            String targetTableName = ctx.targetTable.getTableName();
            final List<String> columns = ctx.targetTable.getColumns();
            final List<String> partitionColumns = ctx.targetTable.getPartitionColumns(); // 获取分区列

            // 检查是否已经存在相同的 SchemaTable
            if (schemaTableMap.containsKey(targetTableName)) {
                // 如果存在，合并列
                SchemaTable existingSchemaTable = schemaTableMap.get(targetTableName);
                existingSchemaTable.setColumns(mergeColumns(existingSchemaTable.getColumns(), columns));
                existingSchemaTable.setPartitionColumns(mergeColumns(existingSchemaTable.getPartitionColumns(), partitionColumns));
                // 如果目标表有分区列信息，使用目标表的分区列
                if (partitionColumns != null && !partitionColumns.isEmpty()) {
                    SchemaTable updatedSchemaTable = new SchemaTable(
                        existingSchemaTable.getSchemaName(),
                        targetTableName,
                        existingSchemaTable.getColumns(),
                        existingSchemaTable.getPartitionColumns()
                    );
                    result.remove(existingSchemaTable);
                    result.add(updatedSchemaTable);
                    schemaTableMap.put(targetTableName, updatedSchemaTable);
                }
            } else {
                // 如果不存在，创建新的 SchemaTable 并添加到集合中
                SchemaTable schemaTable = new SchemaTable(
                    ctx.targetTable.getSchemaName(),
                    targetTableName,
                    columns,
                    partitionColumns  // 包含分区列
                );
                result.add(schemaTable);
                schemaTableMap.put(targetTableName, schemaTable);
            }
        }

        return result;
    }

    /**
     * 获取源表信息（FROM 子句中的表）
     */
    public Set<SchemaTable> getSourceTables() {
        SqlMetadataContext ctx = context.get();
        Set<SchemaTable> result = new HashSet<>();
        Map<String, SchemaTable> schemaTableMap = new HashMap<>();

        // 1. 先收集所有实际表（非别名表）的列信息
        Map<String, Set<String>> realTableColumns = new HashMap<>();
        for (Map.Entry<String, String> aliasEntry : ctx.tableAliases.entrySet()) {
            String alias = aliasEntry.getKey();
            String realTable = aliasEntry.getValue();
            if (ctx.sourceTablesWithColumns.containsKey(alias)) {
                realTableColumns.computeIfAbsent(realTable, k -> new HashSet<>())
                        .addAll(ctx.sourceTablesWithColumns.get(alias));
            }
        }

        // 2. 处理所有表
        for (Map.Entry<String, Set<String>> entry : ctx.sourceTablesWithColumns.entrySet()) {
            String tableName = entry.getKey().trim();
            // 跳过子查询别名和表别名
            if (ctx.subqueryColumns.containsKey(tableName) || ctx.tableAliases.containsKey(tableName)) {
                continue;
            }

            Set<String> allColumns = new HashSet<>(entry.getValue());
            // 合并从别名收集到的列
            if (realTableColumns.containsKey(tableName)) {
                allColumns.addAll(realTableColumns.get(tableName));
            }
            List<String> columns = new ArrayList<>(allColumns);

            // 检查是否已经存在相同的 SchemaTable
            if (schemaTableMap.containsKey(tableName)) {
                // 如果存在，合并列
                SchemaTable existingSchemaTable = schemaTableMap.get(tableName);
                existingSchemaTable.setColumns(mergeColumns(existingSchemaTable.getColumns(), columns));
            } else {
                // 如果不存在，创建新的 SchemaTable 并添加到集合中
                SchemaTable schemaTable = getSchemaTable(tableName, columns,null);
                result.add(schemaTable);
                schemaTableMap.put(tableName, schemaTable);
            }
        }

        return result;
    }

    private List<String> mergeColumns(List<String> existingColumns, List<String> newColumns) {
        // 合并两个列列表并去重
        Set<String> mergedColumnsSet = new HashSet<>(existingColumns);
        mergedColumnsSet.addAll(newColumns);
        return new ArrayList<>(mergedColumnsSet);
    }

    /**
     * 获取目标表信息（INSERT INTO 的表）
     */
    public SchemaTable getTargetTable() {
        return context.get().targetTable;
    }

    public SchemaTable getSchemaTable(String tableName, List<String> columns,List<String> partitionColumns) {
        SqlMetadataContext ctx = context.get();
        // 处理可能的引号
        tableName = tableName.replaceAll("^\"|\"$", "").replaceAll("^`|`$", "");
        if (tableName.contains(".")) {
            String[] parts = tableName.split("\\.");
            if (parts.length == 2) {
                String schemaName = parts[0];
                String tableNamePart = parts[1];
                return new SchemaTable(schemaName, tableNamePart, columns,partitionColumns);
            }
        }
        return new SchemaTable(ctx.defaultSchema,tableName, columns,partitionColumns);
    }

    /**
     * 获取表和列的映射关系
     */
    public Map<String, Set<String>> getTableColumns() {
        return context.get().tableColumns;
    }

    /**
     * 获取表别名到真实表名的映射
     */
    public Map<String, String> getTableAliases() {
        return context.get().tableAliases;
    }

    /**
     * 获取最后一个 SELECT 语句的输出列
     */
    public List<String> getLastSelectColumns() {
        return context.get().lastSelectColumns;
    }
}